use regi::Regi;
use std::io;

fn print_help() {
    eprintln!("Usage: devmem <addr> [width] [val]");
    eprintln!("  <addr>: Physical address to access (in hexadecimal).");
    eprintln!("  [width]: Optional, bits of data width, can be 32 or 64. Default is 32.");
    eprintln!("  [val]: Optional, value to write (in hexadecimal).");
}

fn main() -> io::Result<()> {
    let args: Vec<String> = std::env::args().collect();

    if args.len() == 2 && (args[1] == "-h" || args[1] == "--help") {
        print_help();
        return Ok(());
    }

    if args.len() < 2 {
        print_help();
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,                               
            "Insufficient arguments",
        ));
    }

    let phys_addr = u64::from_str_radix(&args[1].trim_start_matches("0x"), 16)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;

    let width = if args.len() > 2 {
        match args[2].as_str() {
            "32" => 32,
            "64" => 64,
            _ => {
                eprintln!("Invalid width: {}. Allowed bit width only 32 or 64.", args[2]);
                return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid width"));
            }
        }
    } else {
        32
    };

    let size = (width / 8) as usize;
    let mut mapper = Regi::new(phys_addr, size)?;

    if args.len() == 2 || args.len() == 3 {
        // 读操作
        if width == 32 {
            let val32 = mapper.read_u32(0);
            println!("Read u32 at phys_addr {:#x}: {:#x}", phys_addr, val32);
        } else {
            let val64 = mapper.read_u64(0);
            println!("Read u64 at phys_addr {:#x}: {:#x}", phys_addr, val64);
        }
    } else if args.len() == 4 {
        // 写操作
        if width == 32 {
            let value = u32::from_str_radix(&args[3].trim_start_matches("0x"), 16)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
            mapper.write_u32(0, value);
            println!("Wrote u32 {:#x} at phys_addr {:#x}", value, phys_addr);
        } else {
            let value = u64::from_str_radix(&args[3].trim_start_matches("0x"), 16)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
            mapper.write_u64(0, value);
            println!("Wrote u64 {:#x} at phys_addr {:#x}", value, phys_addr);
        }
    } else {
        print_help();
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "Too many arguments",
        ));
    }

    Ok(())
}
